import json
import datetime
from typing import List, Callable, TypeVar, Generator, Tuple, Optional, Union, Any

import peewee
from peewee import fn
from shapely.geometry.base import BaseGeometry

from catalog.models import KNOWN_MODELS
from catalog.models.base_stac_item import BaseSTACItem
from common.geoimage.virtual_raster import VRTSceneCollection

T = TypeVar("T", bound="BaseSTACItem")
IC = TypeVar("IC", bound="ImageCollection")


class ImageCollection(object):
    def __init__(self) -> None:
        self.query = None
        self.lst_images = None
        self.lst_assets = None

    @staticmethod
    def from_collection_name(collection_name: str) -> IC:
        """Create an ImageCollection from a given collection name

        Parameters
        ----------
        collection_name : str
            A known collection name, an example is "COPERNICUS/SENTINEL2/L2A".

        Returns
        -------
        ImageCollection
            An ImageCollection for the given collection.

        Raises
        ------
        ValueError
            Raised when the given `collection_name` does not represent a valid collection in the catalog.
        """

        try:
            model = KNOWN_MODELS[collection_name]
        except KeyError:
            raise ValueError(
                f"Unknown collection {collection_name}, only support {list(KNOWN_MODELS.keys())}"
            )
        else:
            ic = ImageCollection()
            ic.query = model.select()
            return ic

    @staticmethod
    def from_query(query: peewee.Query) -> IC:
        """Create an ImageCollection from a given Query object

        Parameters
        ----------
        query : peewee.Query
            A query object that will be used to search the database and generate images in the ImageCollection object.

        Returns
        -------
        ImageCollection
            An ImageCollection object with images from the given query.
        """
        ic = ImageCollection()
        ic.query = query
        return ic

    @staticmethod
    def from_images(images: List[BaseSTACItem]) -> IC:
        """Create an ImageCollection from a list of images.

        Parameters
        ----------
        images : List[BaseSTACItem]
            A list of images to form the ImageCollection object.
            Each image in the list is an object of the `BaseSTACItem` class.

        Returns
        -------
        ImageCollection
            An ImageCollection object created from the list of images.
        """
        ic = ImageCollection()
        ic.lst_images = images
        return ic

    @staticmethod
    def from_image_collection(image_collection: IC) -> IC:
        """Create an ImageCollection object

        Parameters
        ----------
        image_collection : ImageCollection
            An ImageCollection to copy from.

        Returns
        -------
        ImageCollection
            An ImageCollection copied from the given one.
        """
        ic = ImageCollection()
        ic.query = image_collection.query
        ic.lst_images = image_collection.lst_images
        ic.lst_assets = image_collection.lst_assets
        return ic

    def filter_date(
        self,
        start_datetime: Tuple[
            Union[str, datetime.datetime], Union[str, datetime.datetime]
        ],
        end_datetime: Tuple[
            Union[str, datetime.datetime], Union[str, datetime.datetime]
        ],
    ) -> IC:
        """Create an ImageCollection object by filtering the acquired datetime from current ImageCollection

        Parameters
        ----------
        start_datetime : Tuple[ Union[str, datetime.datetime], Union[str, datetime.datetime] ]
            Start datetime criteria for the image acquired time.
        end_datetime : Tuple[ Union[str, datetime.datetime], Union[str, datetime.datetime] ]
            End datetime criteria for the image acquired time.

        Returns
        -------
        ImageCollection
            An ImageCollection filtered by the acquired datetime.

        Raises
        ------
        ValueError
            Raised when the ImageCollection object is not correctly initialized.
        """
        if self.query is not None:
            new_query = self.query.where(
                self.query.model.start_datetime >= start_datetime
            ).where(self.query.model.end_datetime <= end_datetime)
            return ImageCollection.from_query(new_query)
        elif self.lst_images is not None:
            new_lst_images = []
            for image in self.lst_images:
                if (image.start_datetime >= start_datetime) and (
                    image.end_datetime <= end_datetime
                ):
                    new_lst_images.append(image)
            return ImageCollection.from_images(new_lst_images)
        else:
            raise RuntimeError

    def filter_bounds(self, bbox):
        # TODO add it later to mimic GEE
        raise NotImplementedError

    def filter_aoi(self, aoi: BaseGeometry) -> IC:
        """Create an ImageCollection object by filtering the acquired datetime from current ImageCollection

        Parameters
        ----------
        aoi : BaseGeometry
            The area of interest.

        Returns
        -------
        ImageCollection
            An ImageCollection object filtered by Area-of-Interest.

        Raises
        ------
        RuntimeError
            Raised when the ImageCollection object is not initialized.
        """
        if self.query is not None:
            new_query = self.query.where(
                fn.ST_Intersects(
                    self.query.model.geom, fn.ST_GeomFromGeoJSON(json.dumps(aoi))
                )
            )
            return ImageCollection.from_query(new_query)
        elif self.lst_images is not None:
            new_lst_images = []
            for image in self.lst_images:
                if image.geom.intersect(aoi):
                    new_lst_images.append(image)
            return ImageCollection.from_images(new_lst_images)
        else:
            raise ValueError

    def filter_metadata(self, field: str, operator: str, value: Any) -> IC:  # noqa C901
        """Create an ImageCollection object by filtering from current ImageCollection on a given field

        Parameters
        ----------
        field : str
            A field to be filtered.
        operator : str
            The operator used to perform the filtering. Can be one of :
                - "equals"
                - "less_than"
                - "greater_than"
                - "not_equals"
                - "not_less_than"
                - "not_greater_than"
                - "starts_with"
                - "ends_with"
                - "not_starts_with"
                - "not_ends_with"
                - "contains"
                - "not_contains"
        value : Any
            The value that to be filtered

        Returns
        -------
        I
            [description]

        Raises
        ------
        RuntimeError
            Raised under on of the following circumstances:
                - The given field is not a known field for the collection.
                - The given operator is not one of the known operators.
        NotImplementedError
            We have not implemented the mechanism for image list based ImageCollection yet.
        """
        if self.query is None:
            # only available when it is from a query
            raise NotImplementedError

        KNOWN_OPERATORS = (
            "equals",
            "less_than",
            "greater_than",
            "not_equals",
            "not_less_than",
            "not_greater_than",
            "starts_with",
            "ends_with",
            "not_starts_with",
            "not_ends_with",
            "contains",
            "not_contains",
        )
        if operator not in KNOWN_OPERATORS:
            raise ValueError(f"Unknown operator {operator}")

        known_fields = set(self.query.model._meta.fields().keys()) - set(
            ["start_datetime", "end_datetime", "info", "geom"]
        )
        if field not in known_fields:
            raise ValueError(f"Unknown field {field}")

        if operator == "equals":
            new_query = self.query.where(self.query.model._meta.fields[field] == value)
        elif operator == "less_than":
            new_query = self.query.where(self.query.model._meta.fields[field] < value)
        elif operator == "greater_than":
            new_query = self.query.where(self.query.model._meta.fields[field] > value)
        elif operator == "not_equals":
            new_query = self.query.where(self.query.model._meta.fields[field] != value)
        elif operator == "not_less_than":
            new_query = self.query.where(self.query.model._meta.fields[field] >= value)
        elif operator == "not_greater_than":
            new_query = self.query.where(self.query.model._meta.fields[field] <= value)
        elif operator == "starts_with":
            new_query = self.query.where(
                self.query.model._meta.fields[field].startswith(value)
            )
        elif operator == "ends_with":
            new_query = self.query.where(
                self.query.model._meta.fields[field].endswith(value)
            )
        elif operator == "not_starts_with":
            new_query = self.query.where(
                ~(self.query.model._meta.fields[field].startswith(value))
            )
        elif operator == "not_ends_with":
            new_query = self.query.where(
                ~(self.query.model._meta.fields[field].endswith(value))
            )
        elif operator == "contains":
            new_query = self.query.where(
                self.query.model._meta.fields[field].contains(value)
            )
        elif operator == "not_contains":
            new_query = self.query.where(
                ~(self.query.model._meta.fields[field].contains(value))
            )

        return ImageCollection.from_query(new_query)

    def sort(self, field, ascending=True):
        known_fields = set(self.query.model._meta.fields().keys()) - set(
            ["start_datetime", "end_datetime", "info", "geom"]
        )
        if field not in known_fields:
            raise ValueError(f"Unknown field {field}")

        if self.query is not None:
            if ascending:
                new_query = self.query.order_by(self.query._meta.fields[field])
            else:
                new_query = self.query.order_by(self.query._meta.fields[field].desc())
            return ImageCollection.from_query(new_query)

    def size(self):
        if self.query is not None:
            return self.query.count()
        elif self.lst_images is not None:
            return len(self.lst_images)
        else:
            raise RuntimeError

    def to_list(self, offset=0, count=None):
        if self.query is not None:
            res = []
            for image in self.query.offset(offset).limit(count):
                res.append(image)
            return res
        elif self.lst_images is not None:
            res = []
            for ind, image in enumerate(self.lst_images):
                if ind >= offset:
                    if (count is not None) and (ind - offset <= count):
                        res.append(image)
            return res
        else:
            raise RuntimeError

    def select(self, assets):
        ic = ImageCollection.from_image_collection(self)
        ic.lst_assets = assets
        return ic

    def to_scenes(self, assets: Optional[List[str]] = None) -> VRTSceneCollection:
        if assets is None:
            assets = self.lst_assets
        if assets is None:
            raise RuntimeError("should specify assets before hand")

        res = []
        for image in self.lst_images:
            vrt_scene = image.choose_assets(assets)
            res.append(vrt_scene)
        return VRTSceneCollection(res)

    def group_by(
        self, group_func: Callable[[T], str]
    ) -> Generator[Tuple[str, IC], None, None]:
        groups = {}
        for image in self.query:
            group_key = group_func(image)
            if group_key in groups:
                groups[group_key].append(image)
            else:
                groups[group_key] = [image]
        for key, value in groups.item():
            ic = ImageCollection()
            ic.lst_images = value
            yield key, ic
